"""
QMIX神经网络实现模块

该模块实现了QMIX (Q-value Mixing) 算法的神经网络结构，用于多智能体强化学习。
QMIX是一种值分解方法，能够将团队的联合Q值分解为各个智能体的个体Q值。
主要包含Q网络、经验回放缓冲区和QMIX算法主体三个核心组件。
"""

import torch  # PyTorch深度学习框架
import torch.nn as nn  # 神经网络模块
import torch.optim as optim  # 优化器模块
import numpy as np  # 数值计算库
import random  # 随机数生成
from Utils.tensor_ops import _2tensor  # 张量操作工具
from collections import namedtuple  # 命名元组，用于存储经验回放数据
from Utils.tensor_ops import pt_inf  # 获取PyTorch无穷大值
from torch.distributions.categorical import Categorical  # 分类分布，用于动作采样

# 定义经验回放缓冲区中存储的转换数据结构
# 每个转换包含状态、动作、下一状态、奖励和是否完成标志
Transition = namedtuple('Transition', ('state', 'action', 'next_state', 'reward', 'done'))


class ReplayBuffer:
    """
    经验回放缓冲区类

    用于存储和采样智能体与环境的交互历史数据。
    经验回放是强化学习中的重要技术，可以打破数据相关性，提高样本利用率。
    """

    def __init__(self, capacity):
        """
        初始化回放缓冲区

        Args:
            capacity (int): 缓冲区最大容量，超过后将覆盖最旧的数据
        """
        self.capacity = capacity  # 缓冲区容量
        self.memory = []  # 存储转换数据的列表
        self.position = 0  # 当前写入位置

    def push(self, *args):
        """
        将新的转换数据存入缓冲区

        Args:
            *args: 包含(state, action, next_state, reward, done)的参数
        """
        # 如果缓冲区未满，添加一个空位
        if len(self.memory) < self.capacity:
            self.memory.append(None)
        # 将新数据存储到当前位置
        self.memory[self.position] = Transition(*args)
        # 更新位置，实现循环覆盖
        self.position = (self.position + 1) % self.capacity

    def sample(self, batch_size):
        """
        从缓冲区中随机采样一批数据

        Args:
            batch_size (int): 采样批次大小

        Returns:
            list: 包含batch_size个转换数据的列表
        """
        return random.sample(self.memory, batch_size)

    def __len__(self):
        """
        获取当前缓冲区中的数据数量

        Returns:
            int: 缓冲区中存储的转换数量
        """
        return len(self.memory)


class QNetwork(nn.Module):
    """
    Q网络类

    实现了用于估计动作价值的深度神经网络。
    采用三层全连接结构，用于每个智能体的独立Q值估计。
    """

    def __init__(self, input_dim, output_dim):
        """
        初始化Q网络

        Args:
            input_dim (int): 输入特征维度（状态空间维度）
            output_dim (int): 输出维度（动作空间维度）
        """
        super(QNetwork, self).__init__()
        # 定义三层全连接网络结构
        self.fc1 = nn.Linear(input_dim, 512)  # 第一层：输入层到隐藏层(512单元)
        self.fc2 = nn.Linear(512, 256)       # 第二层：隐藏层到隐藏层(256单元)
        self.fc3 = nn.Linear(256, output_dim) # 第三层：隐藏层到输出层(Q值)

    def forward(self, x):
        """
        前向传播

        Args:
            x (torch.Tensor): 输入状态张量

        Returns:
            torch.Tensor: 各动作的Q值
        """
        # 通过三层全连接层，前两层使用ReLU激活函数
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)  # 输出层不使用激活函数，直接输出Q值
        return x


class QMIX(nn.Module):
    """
    QMIX算法主体类

    实现了QMIX多智能体强化学习算法的核心功能。
    QMIX通过值分解的方式，在保证单调性的前提下学习团队的联合Q值函数。
    每个智能体拥有独立的Q网络，通过训练实现协作。
    """

    def __init__(self, num_agents, state_dim, action_space, device):
        """
        初始化QMIX算法

        Args:
            num_agents (int): 智能体数量
            state_dim (int): 单个智能体的状态空间维度
            action_space (int): 动作空间大小
            device (torch.device): 计算设备(CPU/GPU)
        """
        super().__init__()
        self.num_agents = num_agents  # 智能体数量
        self.state_dim = state_dim    # 状态空间维度
        self.action_space = action_space  # 动作空间大小
        self.device = device          # 计算设备
        self.num_thread = 0           # 线程数量(用于并行处理)

        # 为每个智能体创建独立的Q网络
        # 使用ModuleList管理多个Q网络参数
        self.q_networks = torch.nn.ModuleList(
            modules=[QNetwork(self.state_dim, self.action_space) for _ in range(self.num_agents)])

        # 初始化经验回放缓冲区，容量设为10000
        self.replay_buffer = ReplayBuffer(capacity=10000)

        # 为每个Q网络创建独立的Adam优化器，学习率为0.001
        self.optimizers = [optim.Adam(q_net.parameters(), lr=0.001) for q_net in self.q_networks]
        # 创建优化器字典，便于管理和访问
        self.optimizers_dict = {'optimizer' + str(index + 1): optimizer for index, optimizer in
                                enumerate(self.optimizers)}

    def select_action(self, states, test_mode, avail_act, epsilon=0.1):
        """
        根据当前状态选择动作

        实现了ε-贪心策略结合动作掩码的动作选择机制。
        考虑可用动作约束，使用分类分布进行动作采样。

        Args:
            states (torch.Tensor): 当前状态张量
            test_mode (bool): 是否为测试模式(测试时选择最优动作)
            avail_act (list): 可用动作掩码列表
            epsilon (float): 探索率，默认0.1

        Returns:
            numpy.ndarray: 选择的动作数组
        """
        flag = True  # 标志位，用于处理Q值拼接
        q_value = None  # 存储Q值张量
        self.num_thread = states.shape[0]  # 获取线程数量
        action_list = torch.empty(1, self.num_agents).to(self.device)  # 初始化动作列表

        # 对每个线程进行处理
        for i in range(self.num_thread):
            # 重塑状态张量为(num_agents, state_dim)
            state_thread = states[i].reshape(self.num_agents, self.state_dim)
            # 将可用动作转换为张量
            avail_act_thread = torch.tensor(avail_act[i],dtype=torch.int64).to(self.device)

            # 为每个智能体计算Q值
            for q_network, state in zip(self.q_networks, state_thread):
                # 将状态转换为张量并处理NaN值(替换为0)
                state_tensor = torch.tensor(state, dtype=torch.float32).to(self.device)
                state_tensor = torch.nan_to_num_(state_tensor, 0).to(self.device)

                # 小概率随机探索(1%概率)
                if random.random() < 0.01:
                    action = random.randint(0, self.action_space - 1)
                    q_values = torch.zeros(self.action_space).to(self.device)
                    q_values[action] = 1
                else:
                    # 使用Q网络预测Q值，不计算梯度
                    with torch.no_grad():
                        q_values = q_network(state_tensor).to(self.device)

                # 拼接所有智能体的Q值
                if flag:
                    q_value = q_values.unsqueeze(0)
                    flag = False
                else:
                    q_value = torch.cat((q_value, q_values.unsqueeze(0)))

            # 应用动作掩码：将不可用动作的Q值设为负无穷
            if avail_act is not None:
                q_value = torch.where(avail_act_thread > 0, q_value, -pt_inf()).to(self.device)

            # 使用分类分布进行动作采样
            act_dist = Categorical(logits=q_value.unsqueeze(0))
            if not test_mode:
                # 训练模式：根据概率分布采样
                actions = act_dist.sample()
            else:
                # 测试模式：选择概率最大的动作
                actions = torch.argmax(act_dist.probs, dim=2)

            # 将动作添加到动作列表
            if i == 0:
                action_list = actions
            else:
                action_list = torch.cat((action_list, actions), 0)

            # 重置临时变量
            q_value = None
            flag = True

        return action_list.cpu().numpy()

    def evaluate_actions(self, *args, **kargs):
        """
        评估动作的包装函数

        直接调用_update方法进行动作评估和模型更新。

        Args:
            *args: 位置参数
            **kargs: 关键字参数

        Returns:
            更新后的损失值
        """
        return self._update(*args, **kargs)

    def _update(self, state_batch, action_batch, avail_act_batch, next_state_batch,
                reward_batch, done_batch, gamma, max_grad_norm):
        """
        更新Q网络参数

        使用时序差分(TD)学习算法更新所有智能体的Q网络。
        包含目标网络计算、损失计算、反向传播和梯度裁剪等步骤。

        Args:
            state_batch (torch.Tensor): 当前状态批次
            action_batch (torch.Tensor): 执行的动作批次
            avail_act_batch (torch.Tensor): 可用动作批次
            next_state_batch (torch.Tensor): 下一状态批次
            reward_batch (torch.Tensor): 奖励批次
            done_batch (torch.Tensor): 完成标志批次
            gamma (float): 折扣因子
            max_grad_norm (float): 最大梯度范数，用于梯度裁剪

        Returns:
            float: 平均损失值
        """
        loss_list = []  # 存储每个智能体的损失

        # 对每个智能体独立更新Q网络
        for i in range(self.num_agents):
            # 重塑张量维度：(batch_size, sequence_length, state_dim) -> (batch_size*sequence_length, state_dim)
            state = state_batch[:, i, :, :].reshape(state_batch.shape[0], -1)
            next_state = next_state_batch[:, i, :, :].reshape(next_state_batch.shape[0], -1)

            # 处理NaN值，替换为0
            state = torch.nan_to_num_(state, 0)
            next_state = torch.nan_to_num_(next_state, 0)

            # 提取相关数据
            action = action_batch[:, i].to(torch.int64)  # 转换为整数类型
            avail_act = avail_act_batch[:, i, :]        # 可用动作掩码
            reward = reward_batch[:, i]                 # 奖励值
            done = done_batch[:, i]                     # 完成标志

            # 计算目标Q值（不计算梯度）
            with torch.no_grad():
                # 使用Q网络预测下一状态的Q值
                q_values_next = self.q_networks[i](next_state)
                # 应用动作掩码
                if avail_act is not None:
                    q_values_next = torch.where(avail_act > 0, q_values_next, -pt_inf())
                # 选择最大Q值
                q_values_next = q_values_next.max(1)[0]
                # 计算目标Q值：即时奖励 + 折扣的未来Q值(如果未完成)
                q_targets = reward + gamma * q_values_next * ~done

            # 计算当前状态的Q值
            q_values = self.q_networks[i](state)
            # 计算损失：实际执行的动作的Q值与目标Q值的L1损失
            loss = nn.functional.l1_loss(q_values.gather(1, action.unsqueeze(1)), q_targets.unsqueeze(1))
            loss_list.append(loss.item())

            # 反向传播更新网络参数
            self.optimizers[i].zero_grad()  # 清零梯度
            loss.backward()                  # 反向传播
            # 梯度裁剪，防止梯度爆炸
            nn.utils.clip_grad_norm_(self.parameters(), max_grad_norm)
            self.optimizers[i].step()        # 更新参数

        # 计算平均损失
        mean_loss = np.array(loss_list).sum() / (state_batch.shape[0] * state_batch.shape[1])
        return mean_loss

    def store_transition(self, transition):
        """
        存储转换数据到回放缓冲区

        Args:
            transition (tuple): 包含(state, action, next_state, reward, done)的转换数据
        """
        self.replay_buffer.push(*transition)